use cudarc::driver::{CudaSlice, DriverError};
use glam::Mat4;
use glium::winit::{
    event::{ElementState, KeyEvent, WindowEvent},
    keyboard,
};

use super::{
    camera::Camera,
    cuda_shader::{Device, Shader},
    framebuffer_renderer::CudaRendererOps,
    model::GpuModel,
    tgaimage::{ColorFormat, TGAImage},
};

pub struct CudaRenderer {
    device: Device,
    shader: Shader,

    frame_buffer: Vec<u8>,
    frame_buffer_zero: Vec<u8>,
    frame_buffer_gpu: CudaSlice<u8>,
    z_buffer_zero: Vec<f32>,
    z_buffer_gpu: CudaSlice<f32>,
    frame_buffer_size: (usize, usize),

    models: Vec<GpuModel>,

    camera: Camera,
}

impl CudaRenderer {
    pub fn new(fb_size: (usize, usize), camera: Camera) -> Result<Self, DriverError> {
        let device = Device::new().unwrap();
        let shader = device.load_default_shader().unwrap();

        let frame_buffer = vec![[0u8, 0, 0, 255]; fb_size.0 * fb_size.1].concat();
        let frame_buffer_gpu = device.default_stream().memcpy_stod(&frame_buffer).unwrap();
        let z_buffer_gpu = device
            .default_stream()
            .memcpy_stod(&vec![0f32; fb_size.0 * fb_size.1])
            .unwrap();

        Ok(Self {
            device,
            models: vec![],
            shader,
            frame_buffer,
            frame_buffer_zero: vec![[0u8, 0, 0, 255]; fb_size.0 * fb_size.1].concat(),
            frame_buffer_gpu,
            z_buffer_zero: vec![10000.; fb_size.0 * fb_size.1],
            z_buffer_gpu,
            frame_buffer_size: fb_size,
            camera,
        })
    }

    pub fn add_model(&mut self, model: GpuModel) {
        self.models.push(model);
    }

    pub fn pipeline_render(
        &mut self,
        model_mats: &[Mat4],
        view_mat: &Mat4,
        persp_mat: &Mat4,
        viewport_mat: &Mat4,
    ) {
        self.model_view_transform(model_mats, view_mat);
        self.projection_transform(persp_mat);
        self.ndc_transform();
        self.viewport_transform(viewport_mat);

        self.reset_framebuffer();
        self.rasterization(model_mats, view_mat);
    }

    pub fn device(&self) -> &Device {
        &self.device
    }

    fn reset_framebuffer(&mut self) {
        self.device
            .default_stream()
            .memcpy_htod(&self.frame_buffer_zero, &mut self.frame_buffer_gpu)
            .unwrap();
        self.device
            .default_stream()
            .memcpy_htod(&self.z_buffer_zero, &mut self.z_buffer_gpu)
            .unwrap();
    }

    fn model_view_transform(&mut self, model_mats: &[Mat4], view_mat: &Mat4) {
        for (model, mat) in self.models.iter_mut().zip(model_mats.iter()) {
            let model_view_mat = view_mat.mul_mat4(mat);
            let normal_mat = model_view_mat.inverse().transpose();
            let model_view_mat = self
                .device
                .default_stream()
                .memcpy_stod(&model_view_mat.transpose().to_cols_array())
                .unwrap();
            let normal_mat = self
                .device
                .default_stream()
                .memcpy_stod(&normal_mat.transpose().to_cols_array())
                .unwrap();
            self.shader
                .mat4_transform(
                    model.verts.as_view(),
                    model.verts_camera.as_view_mut(),
                    model_view_mat.as_view(),
                    self.device.default_stream(),
                )
                .unwrap();
            self.shader
                .mat4_transform(
                    model.norms.as_view(),
                    model.norms_camera.as_view_mut(),
                    normal_mat.as_view(),
                    self.device.default_stream(),
                )
                .unwrap();
        }
    }

    fn projection_transform(&mut self, projection_mat: &Mat4) {
        let project_mat = self
            .device
            .default_stream()
            .memcpy_stod(&projection_mat.transpose().to_cols_array())
            .unwrap();
        for model in &mut self.models {
            self.shader
                .mat4_transform(
                    model.verts_camera.as_view(),
                    model.verts_screen.as_view_mut(),
                    project_mat.as_view(),
                    self.device.default_stream(),
                )
                .unwrap()
        }
    }

    fn ndc_transform(&mut self) {
        for model in &mut self.models {
            self.shader
                .vertex_ndc_transform(
                    model.verts_screen.as_view_mut(),
                    self.device.default_stream(),
                )
                .unwrap()
        }
    }

    fn viewport_transform(&mut self, viewport_mat: &Mat4) {
        let viewport_mat = self
            .device
            .default_stream()
            .memcpy_stod(&viewport_mat.transpose().to_cols_array())
            .unwrap();
        for model in &mut self.models {
            self.shader
                .mat4_transform_inplace(
                    model.verts_screen.as_view_mut(),
                    viewport_mat.as_view(),
                    self.device.default_stream(),
                )
                .unwrap()
        }
    }

    fn rasterization(&mut self, model_mats: &[Mat4], view_mat: &Mat4) {
        for (model, model_mat) in self.models.iter_mut().zip(model_mats.iter()) {
            let model_view_mat = view_mat.mul_mat4(model_mat);
            let normal_mat = model_view_mat.inverse().transpose();
            let model_view_mat = self
                .device
                .default_stream()
                .memcpy_stod(&model_view_mat.transpose().to_cols_array())
                .unwrap();
            let normal_mat = self
                .device
                .default_stream()
                .memcpy_stod(&normal_mat.transpose().to_cols_array())
                .unwrap();
            self.shader
                .rasterization(
                    model_view_mat.as_view(),
                    normal_mat.as_view(),
                    model.verts_screen.as_view(),
                    model.facet_vrt.as_view(),
                    model.norms_camera.as_view(),
                    model.facet_vrt.as_view(),
                    model.tex.as_view(),
                    model.facet_tex.as_view(),
                    model.normal_map.as_view(),
                    model.normal_map_size,
                    model.diffuse_map.as_view(),
                    model.diffuse_map_size,
                    self.frame_buffer_gpu.as_view_mut(),
                    self.z_buffer_gpu.as_view_mut(),
                    self.frame_buffer_size,
                    self.device.default_stream(),
                )
                .unwrap();
        }
    }
}

impl CudaRendererOps for CudaRenderer {
    fn gen_framebuffer(&mut self, frame_index: usize, framebuffer_size: (usize, usize)) -> Vec<u8> {
        assert!(self.frame_buffer_size == framebuffer_size);
        let (width, height) = self.frame_buffer_size;

        self.pipeline_render(
            &(0..self.models.len())
                .map(|_| Mat4::IDENTITY)
                .collect::<Vec<_>>(),
            &self.camera.view_mat(),
            &self.camera.perspective_mat(),
            &Camera::view_port_mat(width as f32, height as f32),
        );

        self.device
            .default_stream()
            .memcpy_dtoh(&self.frame_buffer_gpu, &mut self.frame_buffer)
            .unwrap();

        let pixels = self.frame_buffer.clone();

        if self.camera.take_snap && frame_index % 25 == 0 && frame_index < 200 {
            TGAImage::new(width, height, ColorFormat::RGBA)
                .set_data(&pixels)
                .write_tga_file(format!("frame{frame_index}.tga"), true)
                .unwrap();
        }

        pixels
    }

    fn handle_window_event(
        &mut self,
        event: &glium::winit::event::WindowEvent,
        _window: &glium::winit::window::Window,
    ) {
        if let WindowEvent::KeyboardInput {
            event:
                KeyEvent {
                    state: ElementState::Pressed,
                    logical_key: keyboard::Key::Named(key),
                    ..
                },
            ..
        } = event
        {
            self.camera.keyboard_respond(key);
        }
    }
}
